import os

import pandas as pd
import tiktoken

from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
from graphrag.query.indexer_adapters import (
    read_indexer_covariates,
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import (
    LocalSearchMixedContext,
)
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag.vector_stores.lancedb import LanceDBVectorStore

# 输入目录
INPUT_DIR = "./myrag/output"
# LanceDB数据库URI
LANCEDB_URI = f"{INPUT_DIR}/lancedb"

# 社区报告表名
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
# 实体表名
ENTITY_TABLE = "create_final_nodes"
# 实体嵌入表名
ENTITY_EMBEDDING_TABLE = "create_final_entities"
# 关系表名
RELATIONSHIP_TABLE = "create_final_relationships"
# 文本单元表名
TEXT_UNIT_TABLE = "create_final_text_units"
# 社区级别
COMMUNITY_LEVEL = 2

# 读取实体表，获取社区和度数据
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")

# 读取实体数据
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)

# 加载描述嵌入到内存中的LanceDB向量存储中
# 要连接到远程数据库，请指定url和port值。
description_embedding_store = LanceDBVectorStore(
    collection_name="default-entity-description",
)
description_embedding_store.connect(db_uri=LANCEDB_URI)

print(f"实体数量: {len(entity_df)}")
print(entity_df.head())

############
# 读取关系表
relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
# 读取关系数据
relationships = read_indexer_relationships(relationship_df)

print(f"关系数量: {len(relationship_df)}")
print(relationship_df.head())

################

# 读取社区报告表
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
# 读取社区报告数据
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)

print(f"报告记录: {len(report_df)}")
print(report_df.head())

###################

# 读取文本单元表
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
# 读取文本单元数据
text_units = read_indexer_text_units(text_unit_df)

print(f"文本单元记录: {len(text_unit_df)}")
print(text_unit_df.head())

####################




# sd_api_key = "YOUR_DeepSeek_API_KEY"
# openai_api_key = "YOUR_DeepSeek_API_KEY"
# llm_model = "deepseek-chat"
# embedding_model = "text-embedding-3-large"
# llm_api_base = "https://api.deepseek.com"
# embedding_model_api_base = "https://ai.devtool.tech/proxy/v1"

# llm = ChatOpenAI(
#     api_key=sd_api_key,
#     model=llm_model,
#     api_base=llm_api_base,
#     api_type=OpenaiApiType.OpenAI,  
#     max_retries=20,
# )

# token_encoder = tiktoken.get_encoding("cl100k_base")

# text_embedder = OpenAIEmbedding(
#     api_key=openai_api_key,
#     api_base=embedding_model_api_base,
#     api_type=OpenaiApiType.OpenAI,
#     model=embedding_model,
#     deployment_name=embedding_model,
#     max_retries=20,
# )
